#!/usr/bin/env python3
"""
射频腔Ready状态控制器 - 简单版本
用于开发测试，不使用复杂设计模式
"""

import epics
import yaml
import time
import logging
import json
import signal
import sys
from datetime import datetime
from logging.handlers import RotatingFileHandler

# ============================================================================
# 全局变量和常量
# ============================================================================
config = {}
cavities = {}  # {(cm, cav): cavity_info}
logger = None
running = True

# 状态定义
STATE_IDLE = 0
STATE_PREP_DENIED = 1
STATE_LOADING = 2
STATE_READY = 3

STATE_NAMES = {
    0: 'IDLE',
    1: 'PREP_DENIED',
    2: 'LOADING',
    3: 'READY'
}


# ============================================================================
# 配置和初始化
# ============================================================================

def load_config(config_file='config.yaml'):
    """加载配置文件"""
    global config
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        print(f"✓ 配置文件加载成功: {config_file}")
        return True
    except FileNotFoundError:
        print(f"✗ 配置文件不存在: {config_file}")
        return False
    except Exception as e:
        print(f"✗ 配置文件加载失败: {e}")
        return False


def setup_logging():
    """配置日志系统"""
    global logger
    
    log_config = config.get('logging', {})
    log_file = log_config.get('file', 'rf_ready_controller.log')
    log_level = getattr(logging, log_config.get('level', 'INFO'))
    log_format = log_config.get('format', '%(asctime)s - %(levelname)s - %(message)s')
    
    # 创建logger
    logger = logging.getLogger('RF-Ready-Controller')
    logger.setLevel(log_level)
    
    # 文件handler（带轮转）
    max_bytes = log_config.get('max_bytes', 10*1024*1024)
    backup_count = log_config.get('backup_count', 5)
    file_handler = RotatingFileHandler(
        log_file, 
        maxBytes=max_bytes,
        backupCount=backup_count
    )
    file_handler.setLevel(log_level)
    file_handler.setFormatter(logging.Formatter(log_format))
    
    # 控制台handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(log_level)
    console_handler.setFormatter(logging.Formatter(log_format))
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    print(f"✓ 日志系统初始化完成")


def get_cavity_list():
    """生成腔体列表 [(cm_id, cav_id), ...]"""
    hw_config = config['hardware']
    
    # 如果指定了specific_cavities，优先使用
    if 'specific_cavities' in hw_config:
        # 需要解析specific_cavities格式，假设是 ['CM01-01', 'CM01-02']
        result = []
        for item in hw_config['specific_cavities']:
            if '-' in item:
                cm, cav = item.split('-')
                result.append((cm, cav))
            else:
                # 假设是连续编号，需要映射到CM
                cavity_num = int(item)
                cavs_per_cm = len(hw_config.get('cavities_per_cm', ['01','02','03','04','05']))
                cm_num = (cavity_num - 1) // cavs_per_cm + 1
                cav_in_cm = (cavity_num - 1) % cavs_per_cm + 1
                result.append((f'{cm_num:02d}', f'{cav_in_cm:02d}'))
        return result
    
    # 否则根据cryomodules和cavities_per_cm生成
    cms = hw_config.get('cryomodules', ['01'])
    cavs = hw_config.get('cavities_per_cm', ['01'])
    
    result = []
    for cm in cms:
        for cav in cavs:
            result.append((cm, cav))
    
    return result


def build_pv_name(template, cm_id, cav_id):
    """构建PV名称"""
    return template.format(cm=cm_id, cav=cav_id)


# ============================================================================
# 腔体初始化
# ============================================================================

def init_cavity(cm_id, cav_id):
    """初始化单个腔体"""
    cavity_key = (cm_id, cav_id)
    logger.info(f"正在初始化腔体 CM{cm_id}-{cav_id}...")
    
    cavity_info = {
        'cm': cm_id,
        'cav': cav_id,
        'state': STATE_IDLE,
        'pvs': {},
        'rfon_callback_index': None,
        'event_count': 0,
        'last_event_time': time.time(),
        'start_time': time.time()
    }
    
    # 获取配置
    pv_config = config['pv_config']
    timeout = config['runtime']['connection_timeout']
    
    try:
        # 创建所有输入PV
        for pv_key, template in pv_config['input'].items():
            pv_name = build_pv_name(template, cm_id, cav_id)
            pv = epics.PV(pv_name, connection_timeout=timeout, auto_monitor=True)
            cavity_info['pvs'][pv_key] = pv
        
        # 创建所有输出PV
        for pv_key, template in pv_config['output'].items():
            pv_name = build_pv_name(template, cm_id, cav_id)
            pv = epics.PV(pv_name, connection_timeout=timeout, auto_monitor=True)
            cavity_info['pvs'][pv_key] = pv
        
        # 等待连接
        logger.debug(f"等待CM{cm_id}-{cav_id}的PV连接...")
        all_connected = True
        failed_pvs = []
        
        for pv_key, pv in cavity_info['pvs'].items():
            if not pv.wait_for_connection(timeout=timeout):
                all_connected = False
                failed_pvs.append(f"{pv_key}({pv.pvname})")
        
        if not all_connected:
            logger.error(f"CM{cm_id}-{cav_id}: 以下PV连接失败: {', '.join(failed_pvs)}")
            return None
        
        logger.info(f"CM{cm_id}-{cav_id}: 所有PV连接成功 ({len(cavity_info['pvs'])}个)")
        
        # 注册回调
        input_keys = list(pv_config['input'].keys())
        for pv_key in input_keys:
            if pv_key in cavity_info['pvs']:
                pv = cavity_info['pvs'][pv_key]
                pv.add_callback(lambda pvname=None, value=None, **kw: 
                               on_pv_change(cavity_key, pvname, value, **kw))
        
        # 设置初始状态
        set_pv_value(cavity_info, 'ready_all', 0)
        
        logger.info(f"CM{cm_id}-{cav_id} 初始化成功")
        return cavity_info
        
    except Exception as e:
        logger.error(f"CM{cm_id}-{cav_id} 初始化失败: {e}", exc_info=True)
        return None


# ============================================================================
# PV操作辅助函数
# ============================================================================

def get_pv_value(cavity_info, pv_key, default=None):
    """安全获取PV值"""
    pv = cavity_info['pvs'].get(pv_key)
    if not pv or not pv.connected:
        return default
    
    try:
        value = pv.get(use_monitor=True, timeout=0.1)
        return value if value is not None else default
    except Exception as e:
        logger.debug(f"读取PV {pv_key} 失败: {e}")
        return default


def set_pv_value(cavity_info, pv_key, value):
    """安全设置PV值"""
    pv = cavity_info['pvs'].get(pv_key)
    if not pv or not pv.connected:
        return False
    
    try:
        timeout = config['runtime'].get('operation_timeout', 1.0)
        pv.put(value, wait=False, timeout=timeout)
        return True
    except Exception as e:
        logger.debug(f"写入PV {pv_key} 失败: {e}")
        return False


# ============================================================================
# 状态机核心逻辑
# ============================================================================

def on_pv_change(cavity_key, pvname, value, **kwargs):
    """PV变化回调"""
    if not running:
        return
    
    cavity = cavities.get(cavity_key)
    if not cavity:
        return
    
    cavity['event_count'] += 1
    cavity['last_event_time'] = time.time()
    
    try:
        process_state_machine(cavity_key)
    except Exception as e:
        logger.error(f"CM{cavity['cm']}-{cavity['cav']}: 状态机处理异常: {e}", 
                    exc_info=True)


def check_ready_conditions(cavity_info):
    """检查所有就绪条件"""
    thresholds = config['pv_config'].get('ready_thresholds', {})
    target_min = thresholds.get('target_status_min', 0.1)
    trig_min = thresholds.get('trig_status_min', 0.1)
    
    loop_status = get_pv_value(cavity_info, 'loop_status', 0)
    target_status = get_pv_value(cavity_info, 'target_status', 0)
    trig_status = get_pv_value(cavity_info, 'trig_status', 0)
    
    # loop_status应该为1（布尔值），其他两个>阈值
    loop_ok = (loop_status == 1 or loop_status > 0.5)
    target_ok = (target_status > target_min)
    trig_ok = (trig_status > trig_min)
    
    return loop_ok and target_ok and trig_ok


def process_state_machine(cavity_key):
    """处理状态机逻辑"""
    cavity = cavities[cavity_key]
    current_state = cavity['state']
    
    # 读取输入
    enable = get_pv_value(cavity, 'enable', 0)
    epk = get_pv_value(cavity, 'epk_mon', 0.0)
    is_ready_all = check_ready_conditions(cavity)
    
    # 状态机逻辑
    if current_state == STATE_IDLE:
        if enable == 1 or enable > 0.5:  # 兼容布尔和数值
            if abs(epk) > 0.01:  # Epk不为0
                enter_state(cavity_key, STATE_LOADING, "Enable=1且Epk!=0")
            else:
                enter_state(cavity_key, STATE_PREP_DENIED, "Enable=1但Epk=0")
    
    elif current_state == STATE_PREP_DENIED:
        if enable == 0 or enable < 0.5:
            enter_state(cavity_key, STATE_IDLE, "Enable=0")
        elif abs(epk) > 0.01:
            enter_state(cavity_key, STATE_LOADING, "Epk恢复")
    
    elif current_state == STATE_LOADING:
        if enable == 0 or enable < 0.5:
            enter_state(cavity_key, STATE_IDLE, "Enable=0")
        elif is_ready_all:
            enter_state(cavity_key, STATE_READY, "就绪条件满足")
    
    elif current_state == STATE_READY:
        if enable == 0 or enable < 0.5:
            enter_state(cavity_key, STATE_IDLE, "Enable=0")
        elif not is_ready_all:
            enter_state(cavity_key, STATE_LOADING, "就绪条件丢失")


# ============================================================================
# 状态进入动作
# ============================================================================

def enter_state(cavity_key, new_state, reason=""):
    """进入新状态"""
    cavity = cavities[cavity_key]
    old_state = cavity['state']
    
    if old_state == new_state:
        return
    
    cm, cav = cavity['cm'], cavity['cav']
    old_name = STATE_NAMES[old_state]
    new_name = STATE_NAMES[new_state]
    
    logger.info(f"CM{cm}-{cav}: {old_name} -> {new_name} ({reason})")
    
    cavity['state'] = new_state
    
    # 执行状态进入动作
    if new_state == STATE_IDLE:
        set_pv_value(cavity, 'ready_all', 0)
        stop_rfon_monitor(cavity_key)
    
    elif new_state == STATE_PREP_DENIED:
        set_pv_value(cavity, 'ready_all', 0)
        set_pv_value(cavity, 'rfon', 0)
        start_rfon_monitor(cavity_key)
    
    elif new_state == STATE_LOADING:
        set_pv_value(cavity, 'ready_all', 0)
        stop_rfon_monitor(cavity_key)
        
        # 复制监测值到设定值
        epk_val = get_pv_value(cavity, 'epk_mon')
        phase_val = get_pv_value(cavity, 'phase_mon')
        
        if epk_val is not None:
            set_pv_value(cavity, 'epk_set', epk_val)
            logger.debug(f"CM{cm}-{cav}: Epk设定 = {epk_val}")
        
        if phase_val is not None:
            set_pv_value(cavity, 'phase_set', phase_val)
            logger.debug(f"CM{cm}-{cav}: Phase设定 = {phase_val}")
    
    elif new_state == STATE_READY:
        set_pv_value(cavity, 'ready_all', 1)


# ============================================================================
# RFon监控
# ============================================================================

def force_rfon_off(cavity_key, value=None, **kwargs):
    """强制关闭RFon的回调"""
    if value is not None and (value == 1 or value > 0.5):
        cavity = cavities[cavity_key]
        cm, cav = cavity['cm'], cavity['cav']
        set_pv_value(cavity, 'rfon', 0)
        logger.warning(f"CM{cm}-{cav}: 在PREP_DENIED状态强制关闭RFon")


def start_rfon_monitor(cavity_key):
    """启动RFon监控"""
    cavity = cavities[cavity_key]
    
    if cavity['rfon_callback_index'] is None:
        pv = cavity['pvs'].get('rfon')
        if pv:
            idx = pv.add_callback(lambda **kw: force_rfon_off(cavity_key, **kw))
            cavity['rfon_callback_index'] = idx
            logger.debug(f"CM{cavity['cm']}-{cavity['cav']}: RFon监控已启动")


def stop_rfon_monitor(cavity_key):
    """停止RFon监控"""
    cavity = cavities[cavity_key]
    
    if cavity['rfon_callback_index'] is not None:
        pv = cavity['pvs'].get('rfon')
        if pv:
            try:
                pv.remove_callback(cavity['rfon_callback_index'])
            except:
                pass
            cavity['rfon_callback_index'] = None
            logger.debug(f"CM{cavity['cm']}-{cavity['cav']}: RFon监控已停止")


# ============================================================================
# 监控和输出
# ============================================================================

def export_status():
    """导出状态到JSON文件"""
    status = {
        'timestamp': datetime.now().isoformat(),
        'total_cavities': len(cavities),
        'cavities': []
    }
    
    # 统计各状态数量
    state_counts = {name: 0 for name in STATE_NAMES.values()}
    
    for (cm, cav), cavity in sorted(cavities.items()):
        state_name = STATE_NAMES[cavity['state']]
        state_counts[state_name] += 1
        
        connected = sum(1 for pv in cavity['pvs'].values() if pv.connected)
        total = len(cavity['pvs'])
        
        uptime = time.time() - cavity['start_time']
        idle_time = time.time() - cavity['last_event_time']
        
        cavity_status = {
            'cm': cm,
            'cav': cav,
            'state': state_name,
            'event_count': cavity['event_count'],
            'events_per_minute': round(cavity['event_count'] / uptime * 60, 2) if uptime > 0 else 0,
            'connected_pvs': connected,
            'total_pvs': total,
            'uptime_seconds': round(uptime, 1),
            'last_event_seconds_ago': round(idle_time, 1)
        }
        status['cavities'].append(cavity_status)
    
    status['state_summary'] = state_counts
    
    # 写入文件
    output_file = config['runtime'].get('status_output', 'rf_ready_status.json')
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(status, f, indent=2, ensure_ascii=False)
    except Exception as e:
        logger.error(f"导出状态失败: {e}")


def print_summary():
    """打印状态摘要"""
    state_counts = {name: 0 for name in STATE_NAMES.values()}
    
    for cavity in cavities.values():
        state_name = STATE_NAMES[cavity['state']]
        state_counts[state_name] += 1
    
    summary_parts = [f"{state}:{count}" for state, count in sorted(state_counts.items())]
    summary = ' | '.join(summary_parts)
    
    logger.info(f"状态摘要: {summary}")


# ============================================================================
# 信号处理
# ============================================================================

def signal_handler(signum, frame):
    """信号处理器"""
    global running
    logger.info(f"收到信号 {signum}，准备关闭...")
    running = False


# ============================================================================
# 清理函数
# ============================================================================

def cleanup():
    """清理资源"""
    logger.info("正在清理资源...")
    
    for cavity_key, cavity in cavities.items():
        try:
            cm, cav = cavity['cm'], cavity['cav']
            logger.debug(f"清理 CM{cm}-{cav}...")
            
            # 关闭输出
            set_pv_value(cavity, 'ready_all', 0)
            
            # 停止监控
            stop_rfon_monitor(cavity_key)
            
            # 移除所有回调
            for pv in cavity['pvs'].values():
                try:
                    pv.clear_callbacks()
                except:
                    pass
        except Exception as e:
            logger.error(f"清理 {cavity_key} 时出错: {e}")
    
    logger.info("清理完成")


# ============================================================================
# 主程序
# ============================================================================

def main():
    """主函数"""
    global running
    
    print("=" * 70)
    print("射频腔Ready状态控制器 - 简单开发版本")
    print("=" * 70)
    print()
    
    # 加载配置
    if not load_config():
        return 1
    
    # 配置日志
    setup_logging()
    logger.info("="*50)
    logger.info("射频腔Ready状态控制器启动")
    logger.info("="*50)
    
    # 注册信号处理
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    # 获取腔体列表
    cavity_list = get_cavity_list()
    logger.info(f"需要监控 {len(cavity_list)} 个腔体")
    
    # 初始化所有腔体
    print(f"\n正在初始化 {len(cavity_list)} 个腔体...")
    print("-" * 70)
    
    success_count = 0
    skip_failed = config['runtime'].get('skip_failed_cavities', True)
    
    for cm_id, cav_id in cavity_list:
        cavity_info = init_cavity(cm_id, cav_id)
        if cavity_info:
            cavities[(cm_id, cav_id)] = cavity_info
            success_count += 1
            print(f"  ✓ CM{cm_id}-{cav_id} 初始化成功")
        else:
            print(f"  ✗ CM{cm_id}-{cav_id} 初始化失败")
            if not skip_failed:
                logger.critical("因初始化失败退出（skip_failed_cavities=false）")
                return 1
    
    if success_count == 0:
        logger.critical("没有任何腔体初始化成功，退出")
        return 1
    
    logger.info(f"成功初始化 {success_count}/{len(cavity_list)} 个腔体")
    print()
    print("=" * 70)
    print(f"✓ {success_count}/{len(cavity_list)} 个腔体正在运行")
    print()
    print("控制器运行中，按 Ctrl+C 停止")
    print("-" * 70)
    
    # 主循环
    status_interval = config['runtime'].get('status_interval', 60)
    last_status_time = time.time()
    loop_interval = 5
    
    try:
        while running:
            time.sleep(loop_interval)
            
            # 定期导出状态
            if time.time() - last_status_time >= status_interval:
                print_summary()
                export_status()
                last_status_time = time.time()
    
    except KeyboardInterrupt:
        logger.info("收到键盘中断")
    
    except Exception as e:
        logger.critical(f"主循环异常: {e}", exc_info=True)
    
    finally:
        print("\n")
        logger.info("正在关闭系统...")
        cleanup()
        export_status()  # 最后导出一次状态
        logger.info("系统已关闭")
        print("✓ 系统已安全关闭")
    
    return 0


if __name__ == '__main__':
    sys.exit(main())